
////////////////////////////////////////////////////////////////////////
//
//   Description:
//     Core of Kernel Module for Resource Container
//
////////////////////////////////////////////////////////////////////////

#include "blockmma.h"
#include <asm/segment.h>
#include <linux/slab.h>
#include <linux/kernel.h>
#include <linux/errno.h>
#include <linux/mm.h>
#include <linux/fs.h>
#include <linux/miscdevice.h>
#include <linux/module.h>
#include <linux/moduleparam.h>
#include <linux/poll.h>
#include <linux/mutex.h>
#include <linux/sched.h>
#include <linux/mutex.h>
#include "core.h"

static int DEFAULT_SIZE = 128 * 128 * sizeof(float);

struct task_node {
    struct task_node* next;
    volatile int finish;
    volatile int exec;
    volatile int clean;
    struct blockmma_cmd* cmd;
    float* a;
    float* b;
    float* c;
    float* user_c;
    int pid;
};

static struct mutex taskid_mutex;
static struct mutex linklist_mutex;
static long taskid_gen = 0;

static volatile struct task_node* head = NULL;

extern struct miscdevice blockmma_dev;

void get_matrix(float* matrix, unsigned long long user_start_addr, unsigned long long tile, unsigned long long row_size){
    float* matrix_start = matrix;
    unsigned long long addr_start = user_start_addr;
    int i;
    for(i=0;i<tile;i++){
        copy_from_user(matrix_start, addr_start, tile * sizeof(float));
        // TODO:check it carefully
        matrix_start += tile;
        addr_start += sizeof(float) * row_size;
    }
}

void set_matrix(float* matrix, unsigned long long user_start_addr, unsigned long long tile, unsigned long long row_size){
    // float* matrix_temp = (float*)kzalloc(sizeof(float) * tile * tile, GFP_KERNEL);
    float* matrix_start = matrix;
    unsigned long long addr_start = user_start_addr;
    int i , j;

    for(i=0;i<tile;i++){
        // copy_from_user(matrix_temp, addr_start, tile * sizeof(float));
        // for(j=0;j<tile;j++){
            // matrix_temp[j] = matrix_start[j] + matrix_temp[j];
        // }
        copy_to_user(addr_start, matrix_start, tile * sizeof(float));
        // TODO:check it carefully
        // matrix_temp += tile;
        matrix_start += tile;
        addr_start += sizeof(float) * row_size;
    }
    // kfree(matrix_temp);
}

void insert_linklist(struct task_node* node_in){
    struct task_node* new_node = (struct task_node*)kmalloc(sizeof(struct task_node), GFP_KERNEL);

    new_node->next = NULL;
    new_node->finish = 0;
    new_node->exec = 0;
    new_node->clean = 0;
    new_node->pid = node_in->pid;
    new_node->cmd = (struct blockmma_cmd*)kmalloc(sizeof(struct blockmma_cmd), GFP_KERNEL);
    memcpy(new_node->cmd, node_in->cmd, sizeof(struct blockmma_cmd));

    long num_row = new_node->cmd->tile;
    long total_size = sizeof(float) * num_row * num_row;

    float* a = (float*)kzalloc(total_size, GFP_KERNEL);
    float* b = (float*)kzalloc(total_size, GFP_KERNEL);
    float* c = (float*)kzalloc(total_size, GFP_KERNEL);

    // printk(KERN_EMERG "[debug] a is %d b is %d c is %d\n", a, b, c);

    get_matrix(a, new_node->cmd->a, new_node->cmd->tile, new_node->cmd->m);
    get_matrix(b, new_node->cmd->b, new_node->cmd->tile, new_node->cmd->n);
    get_matrix(c, new_node->cmd->c, new_node->cmd->tile, new_node->cmd->k);
    
    // copy_from_user(a, new_node->cmd->a, total_size);
    // copy_from_user(b, new_node->cmd->b, total_size);
    // copy_from_user(c, new_node->cmd->c, total_size);

    new_node->a = a;
    new_node->b = b;
    new_node->c = c;

    new_node->user_c = node_in->cmd->c;

    new_node->next = head;
    head = new_node;
}

int remove_linklist(int pid){
    struct task_node* p = head;
    struct task_node* pre = head;
    while(p){
        if(p->pid == pid){
            break;
        }
        pre = p;
        p = p->next;
    }
    if(p == NULL){
        return -1;
    }
    if(p == head){
        head = NULL;
    }else{
        pre->next = p->next;
    }

    // free p
    kfree(p->a);
    kfree(p->b);
    kfree(p->c);
    kfree(p->cmd);
    kfree(p);

    return 0;

}

/**
 * Enqueue a task for the caller/accelerator to perform computation.
 */
long blockmma_send_task(struct blockmma_cmd __user *user_cmd)
{
    struct blockmma_cmd cmd;
    struct task_node user_task;
    if (copy_from_user(&cmd, user_cmd, sizeof(cmd)))
    {
        return -1;
    }
    struct task_node* p = head;
    mutex_lock(&linklist_mutex);
    while(p){
        if(p->pid == current->pid){
            if(p->finish == 1 && p->clean == 0){
                printk(KERN_EMERG "[sync] writeback a result\n");
                // copy_to_user(p->cmd->c, p->c, DEFAULT_SIZE);
                set_matrix(p->c, p->cmd->c, p->cmd->tile, p->cmd->k);
                p->clean = 1;
                kfree(p->a);
                kfree(p->b);
                kfree(p->c);
            }
            if(p->user_c == cmd.c && p->clean == 0){
                mutex_unlock(&linklist_mutex);
                return -1;
            }
        }
        p=p->next;
    }
    mutex_unlock(&linklist_mutex);

    mutex_lock(&taskid_mutex);
    cmd.tid = taskid_gen++;
    mutex_unlock(&taskid_mutex);
    
    user_task.cmd = &cmd;
    user_task.pid = current->pid;
    mutex_lock(&linklist_mutex);
    insert_linklist(&user_task);
    mutex_unlock(&linklist_mutex);
    printk(KERN_EMERG "task sent! task tid: %d  head's pid : %d\n", user_task.cmd->tid, head->pid);
    return 0;
}

/**
 * Return until all outstanding tasks for the calling process are done
 */
int blockmma_sync(struct blockmma_cmd __user *user_cmd)
{
    struct task_node* p = head;
    struct blockmma_cmd cmd;
    if (copy_from_user(&cmd, user_cmd, sizeof(cmd)))
    {
        return -1;
    }

    // tranverse until all is finish for this pid
    p = head;
    
    mutex_lock(&linklist_mutex);
    while(p){
        if(p->pid == current->pid){
            if(p->finish == 0){
                mutex_unlock(&linklist_mutex);
                return -1;
            }else if(p->clean == 0){
                printk(KERN_EMERG "[sync] writeback a result\n");
                // copy_to_user(p->cmd->c, p->c, DEFAULT_SIZE);
                set_matrix(p->c, p->cmd->c, p->cmd->tile, p->cmd->k);
                p->clean = 1;
                kfree(p->a);
                kfree(p->b);
                kfree(p->c);
            }
        }
        p = p->next;
    }
    // free link list
    // while(remove_linklist(current->pid) == 0){}

    mutex_unlock(&linklist_mutex);
    
    return 0;
}

/**
 * Return the task id of a task to the caller/accelerator to perform computation.
 */
int blockmma_get_task(struct blockmma_hardware_cmd __user *user_cmd)
{
    struct task_node* p = head;
    struct blockmma_hardware_cmd cmd;
    if (copy_from_user(&cmd, user_cmd, sizeof(cmd)))
    {
        return -1;
    }
    mutex_lock(&linklist_mutex);
    while(p){
        if(!(p->finish) && !(p->exec)){
            printk(KERN_EMERG "[task] found a task , task id is %d task's pid is %d\n", p->cmd->tid, p->pid);
            copy_to_user(cmd.a, p->a, DEFAULT_SIZE);
            copy_to_user(cmd.b, p->b, DEFAULT_SIZE);
            copy_to_user(cmd.c, p->c, DEFAULT_SIZE);
            break;
        }
        p = p->next;
    }
    if(p == NULL){
        // do not find a task
        mutex_unlock(&linklist_mutex);
        return -1;
    }
    // we find a task and return the task id 
    p->exec = 1;
    int result = p->cmd->tid;
    mutex_unlock(&linklist_mutex);
    return result;
}


/**
 * Return until the task specified in the command is done.
 */
int blockmma_comp(struct blockmma_hardware_cmd __user *user_cmd)
{
    struct task_node* p = head;
    struct blockmma_hardware_cmd cmd;
    if (copy_from_user(&cmd, user_cmd, sizeof(cmd)))
    {
        return -1;
    }
    mutex_lock(&linklist_mutex);
    // printk(KERN_EMERG "[task] complete a task , task id is %d\n", cmd.tid);
    while(p){
        if(p->cmd->tid == cmd.tid){
            // this task is down by acc
            printk(KERN_EMERG "[task] complete and finish a task cmd.c is %d p->c is %d\n", cmd.c, p->c);
            copy_from_user(p->c, cmd.c, DEFAULT_SIZE);
            p->finish = 1;
            // printk(KERN_EMERG "[task] head's finish is %d\n", head->finish);
            break;
        }
        p = p->next;
    }
    mutex_unlock(&linklist_mutex);
    return 0;
}

/*
 * Tell us who wrote the module
 */
int blockmma_author(struct blockmma_hardware_cmd __user *user_cmd)
{
    struct blockmma_hardware_cmd cmd;
    char authors[] = "Yu-Chia Liu (yliu719), 987654321 and Hung-Wei Tseng (htseng), 123456789";
    if (copy_from_user(&cmd, user_cmd, sizeof(cmd)))
    {
        return -1;
    }
    copy_to_user((void *)cmd.a, (void *)authors, sizeof(authors));
    return 0;
}

int blockmma_init(void)
{
    int ret =0;
    if ((ret = misc_register(&blockmma_dev)))
    {
        printk(KERN_ERR "Unable to register \"blockmma\" misc device\n");
        return ret;
    }
    printk("BlockMMA kernel module installed\n");

    mutex_init(&taskid_mutex);
    mutex_init(&linklist_mutex);
    return ret;
}

void blockmma_exit(void)
{
    printk("BlockMMA removed\n");
    misc_deregister(&blockmma_dev);
}

